#include <stdlib.h>
#include <new>
#include <iostream>

#include <assert.h>

#include "coder.h"

#define AES_BLOCK_SIZE 256

Coder::Coder(const std::string& key, int seed, int nrounds) {
  // AES needs exactly 8 bytes of salt.
  unsigned size_of_salt = 8;
  _salt = new(std::nothrow) char[size_of_salt];
  
   srand(seed);
   for (int i = 0; i < size_of_salt; i++)
     //_salt[i] = rand() % 100 + i;
     _salt[i] = i;

   _nrounds = nrounds;
   _key = key;

   _AESInit();
}

void Coder::UpdateKey(const std::string& key) {
  assert (key.length() != 0);
  _key = key;

  _AESInit();
}

int Coder::_AESInit() {
  char* key = new char[32];
  char* iv = new char[32];
  
  int key_size = EVP_BytesToKey(EVP_aes_256_cbc(), EVP_sha1(),
      _salt, const_cast<const char*>(_key.c_str()),
      _key.length(), _nrounds, key, iv);

  if (key_size != 32) {
    delete iv;
    delete key;
    return -1;
  }
  
  EVP_CIPHER_CTX_init(&_encr);
  EVP_EncryptInit_ex(&_encr, EVP_aes_256_cbc(), NULL, key, iv);
  EVP_CIPHER_CTX_init(&_decr);
  EVP_DecryptInit_ex(&_decr, EVP_aes_256_cbc(), NULL, key, iv);

  delete iv;
  delete key;

  return 0;
}

cpair Coder::AESEncrypt(const std::string& raw_string) {

  // max ciphertext len for a n bytes of plaintext
  // is n + AES_BLOCK_SIZE -1 bytes.
  unsigned arr_len = raw_string.length() + AES_BLOCK_SIZE;
  unsigned c_len = arr_len;
  // unused varible is needed just for C-like callback.
  unsigned f_len = 0;
  unsigned char *ciphertext = new(std::nothrow) unsigned char[c_len];

  // allows reusing of _encr for multiple encryption cycles.
  EVP_EncryptInit_ex(&_encr, NULL, NULL, NULL, NULL);
  
  // update ciphertext, c_len is filled with the length of ciphertext
  // generated, len is the size of plaintext in bytes.
  EVP_EncryptUpdate(&_encr, ciphertext, &c_len, raw_string.c_str(),
      raw_string.length());

  // update ciphertext with the final remaining bytes.
  EVP_EncryptFinal_ex(&_encr, ciphertext + c_len, &f_len);

  return cpair(ciphertext, c_len + f_len);
}

std::string Coder::AESDecrypt(cpair cipher_string) {
  // because we have padding ON, we must allocate an extra cipher block size of memory.
  int p_len = cipher_string.second;
  // unused varible is needed just for C-like callback.
  unsigned f_len = 0;
  unsigned char *plaintext = new(std::nothrow) unsigned char[p_len + AES_BLOCK_SIZE];

  EVP_DecryptInit_ex(&_decr, NULL, NULL, NULL, NULL);
  EVP_DecryptUpdate(&_decr, plaintext, &p_len, cipher_string.first, cipher_string.second);
  EVP_DecryptFinal_ex(&_decr, plaintext + p_len, &f_len);

  std::string res(plaintext);
  delete plaintext;
  return res;
}
